#!/usr/bin/env python3
'''
todo
完成适配，参考连接：https://github.com/highcansavci/reinforcement-learning-a2c/blob/master/model/a2c_model.py#L35

训练记录：
改为在主力机上训练
20250113:训练分数-3.875，测试分数没有输出，爆显存了，降低batch_size
20250114：降低显存占用，训练分数-3.693，测试分数-1.0，加入lstm看起来确实难以训练
20250115: Learning rate: 0.0007，训练分数-3.75分，测试分数-1.0分，修复保存最好模型被转换为int导致模型被覆盖的问题，继续训练一天
20250116：修复保存最好模型会被覆盖的问题，今天暂停训练一天，继续训练
20250117:Learning rate: 0.0007，训练分数-3.717，测试分数-1.0分，继续训练
20250119: 暂停训练一天，继续训练，明天继续训练
20250120：等待breakout训练一天后，继续训练，明天继续训练
20250121: 暂停一天训练，明天继续训练
20250122: 暂停一天训练，明天继续训练，明天继续训练
20250123： Learning rate: 0.0007，训练分数-3.740，测试分数-1.0分，停止训练
'''
import gymnasium as gym
import ptan
import numpy as np
import argparse
from tensorboardX import SummaryWriter
import os

import torch
import torch.nn as nn
import torch.nn.utils as nn_utils
import torch.nn.functional as F
import torch.optim as optim

from typing import Any
from lib import common
from collections import deque
import ale_py

gym.register_envs(ale_py)
GAMMA = 0.99
GAE_LAMBDA = 0.9
VF_COEF = 0.5
ENT_COEF = 0.01
LEARNING_RATE = 7e-4
ENTROPY_BETA = 0.01
BATCH_SIZE = 96
NUM_ENVS = 50

REWARD_STEPS = 4
CLIP_GRAD = 0.5

SAVE_ITERS = 100

class TransposeObservation(gym.ObservationWrapper):
    def __init__(self, env=None):
        super(TransposeObservation, self).__init__(env)

    def observation(self, observation):
        # 将观察从 (H, W, C) 转换为 (C, H, W)
        return observation.transpose(2, 0, 1)


class FireResetEnv(gym.Wrapper):
    def __init__(self, env=None):
        """For environments where the user need to press FIRE for the game to start."""
        super(FireResetEnv, self).__init__(env)
        # 以下可知，一些游戏存在FIRE的动作，并且存在FIRE动作的游戏其游戏动作执行有三个以上
        assert env.unwrapped.get_action_meanings()[1] == 'FIRE'
        assert len(env.unwrapped.get_action_meanings()) >= 3

    def step(self, action):
        return self.env.step(action)

    def reset(self, seed: int | None = None, options: dict[str, Any] | None = None):
        # 这里之所以尝试重置后尝试各种动作，是因为不知道哪个是FIRE，继续游戏，所以一个一个尝试
        # 如果不小心游戏结束了，则继续重置
        # 假设游戏继续游戏的按钮在前3
        self.env.reset(seed=seed, options=options)
        obs, _, done, _, info = self.env.step(1)
        if done:
            self.env.reset(seed=seed, options=options)
        obs, _, done, _, info = self.env.step(2)
        if done:
            self.env.reset(seed=seed, options=options)
        return obs, info

class AtariA2C(nn.Module):
    def __init__(self, input_shape, n_actions, input_size=1):
        super(AtariA2C, self).__init__()

        obs_action = input_shape

        self.conv = nn.Sequential(
            nn.Conv2d(obs_action[0], 32, kernel_size=3, stride=2, padding=1),
            nn.BatchNorm2d(32),
            nn.ELU(True),
            nn.Conv2d(32, 32, kernel_size=3, stride=2, padding=1),
            nn.BatchNorm2d(32),
            nn.ELU(True),
            nn.Conv2d(32, 32, kernel_size=3, stride=2, padding=1),
            nn.BatchNorm2d(32),
            nn.ELU(True),
            nn.Conv2d(32, 32, kernel_size=3, stride=2, padding=1),
            nn.BatchNorm2d(32),
            nn.ELU(True),
            nn.Conv2d(32, 32, kernel_size=3, stride=1, padding=1),
            nn.BatchNorm2d(32),
            nn.ELU(True)
        )

        self.lstm = nn.LSTM(num_layers=1, input_size=input_size, hidden_size=256, batch_first=True)
        self.policy = nn.Sequential(
            nn.Linear(256, 256),
            nn.ReLU(),
            nn.Linear(256, n_actions)
        )

        self.value = nn.Sequential(
            nn.Linear(256, 256),
            nn.ReLU(),
            nn.Linear(256, 1)
        )

        self.input_size = input_size

    def forward(self, x):
        inputs, (hx, cx) = x
        fx = inputs.float() / 255.0
        conv_out = self.conv(fx).view(fx.size()[0], -1, self.input_size)
        lstm_out, (hx, cx) = self.lstm(conv_out, (hx, cx))
        x = lstm_out[:, -1, :]
        return self.policy(x), self.value(x), (hx, cx)


def unpack_batch(batch, net, device='cpu'):
    """
    Convert batch into training tensors
    :param batch:
    :param net:
    :return: states variable, actions tensor, reference values variable
    """
    states = []
    actions = []
    rewards = []
    not_done_idx = [] # 非结束的游戏数据索引，该索引记录对应batch，states，actions，rewards
    last_states = [] # 记录采样中的执行动作后的状态，仅记录游戏非结束状态下的索引
    for idx, exp in enumerate(batch):
        states.append(np.asarray(exp.state))
        actions.append(int(exp.action))
        rewards.append(exp.reward)
        if exp.last_state is not None:
            not_done_idx.append(idx)
            last_states.append(np.asarray(exp.last_state))
    states_v = torch.FloatTensor(np.asarray(states)).to(device)
    actions_t = torch.LongTensor(actions).to(device)
    rewards_np = np.array(rewards, dtype=np.float32)
    if not_done_idx:
        last_states_v = torch.FloatTensor(np.asarray(last_states)).to(device)
        last_vals_v = net(last_states_v)[1]
        last_vals_np = last_vals_v.data.cpu().numpy()[:, 0]
        rewards_np[not_done_idx] += GAMMA ** REWARD_STEPS * last_vals_np

    ref_vals_v = torch.FloatTensor(rewards_np).to(device)   
    return states_v, actions_t, ref_vals_v


class BreakoutPenaltyWrapper(gym.Wrapper):
    def __init__(self, env, frame_penalty=-0.1, life_loss_penalty=-10):
        super(BreakoutPenaltyWrapper, self).__init__(env)
        self.frame_penalty = frame_penalty
        self.life_loss_penalty = life_loss_penalty
        self.previous_lives = 0

    def reset(self, **kwargs):
        obs, info = self.env.reset(**kwargs)
        self.previous_lives = info.get('lives', 0)  # 初始生命值
        return obs, info

    def step(self, action):
        obs, reward, done, truncated, info = self.env.step(action)
        
        # 每帧惩罚
        # reward += self.frame_penalty
        
        # 处理生命减少时的惩罚
        current_lives = info.get('lives', self.previous_lives)
        if current_lives < self.previous_lives:
            reward += self.life_loss_penalty
            self.previous_lives = current_lives
        
        return obs, np.sign(reward), done, truncated, info


def wrap_dqn(env, stack_frames=4, episodic_life=True, reward_clipping=True):
    if episodic_life:
        # 将多条生命的游戏模拟成单条生命ActorCriticAgent
        env = ptan.common.wrappers.EpisodicLifeEnv(env)
    # 增强初始化
    env = ptan.common.wrappers.NoopResetEnv(env, noop_max=30)

    if 'FIRE' in env.unwrapped.get_action_meanings():
        env = FireResetEnv(env)
    env = ptan.common.wrappers.ProcessFrame84(env)
    env = ptan.common.wrappers.ImageToPyTorch(env)
    env = ptan.common.wrappers.FrameStack(env, stack_frames)
    env = BreakoutPenaltyWrapper(env)
    return env


class DQNLstmAgent(ptan.agent.BaseAgent):

    def __init__(self, dqn_model, num_envs, apply_softmax=False, action_selector=ptan.actions.ProbabilityActionSelector(), device="cpu", preprocessor=ptan.agent.default_states_preprocessor):
        '''
        param dqn_model: dqn网络模型，训练的的网络
        param action_selector: 动作选择器
        '''

        self.dqn_model = dqn_model
        self.action_selector = action_selector
        self.preprocessor = preprocessor
        self.device = device
        self.num_envs = num_envs
        self.hx, self.cx = self.initial_state()  # 初始lstm状态
        self.apply_softmax = apply_softmax


    def initial_state(self):
        """
        Should create initial empty state for the agent. It will be called for the start of the episode
        :return: Anything agent want to remember
        """
        return (
            np.zeros((self.dqn_model.lstm.num_layers, self.num_envs, self.dqn_model.lstm.hidden_size), dtype=np.float32),
            np.zeros((self.dqn_model.lstm.num_layers, self.num_envs, self.dqn_model.lstm.hidden_size), dtype=np.float32)
        )
    

    @torch.no_grad()
    def __call__(self, states, agent_states=None):
        # 如果定义了预处理器，则将state进行预处理
        if self.preprocessor is not None:
            states = self.preprocessor(states)
            hx, cx = self.preprocessor(agent_states[0]), self.preprocessor(agent_states[1])
            if torch.is_tensor(states):
                states = states.to(self.device)
                hx = hx.to(self.device)
                cx = cx.to(self.device)
        # 用传入的模型计算预测q值或者动作概率
        logit, value, (hx, cx) = self.dqn_model((states, (hx, cx)))
        # 如果需要使用softmax计算
        if self.apply_softmax:
            probs_v = F.softmax(logit, dim=1)
        value = value.data.cpu().numpy()
        probs = probs_v.data.cpu().numpy()
        actions = self.action_selector(probs)
        hx = hx.cpu().numpy()
        cx = cx.cpu().numpy()
        return actions, value, (hx, cx)

    # 增加保存和加载next_lstm_state的代码
    def save_state_dict(self, checkpoints):
        checkpoints['hx'] = self.hx
        checkpoints['cx'] = self.cx

    def load_state_dict(self, checkpoints):
        self.hx = checkpoints['hx']
        self.cx = checkpoints['cx']


@torch.no_grad()
def test_model(env, net, device, episodes=5):
    total_reward = 0.0
    for _ in range(episodes):
        noop_action_count = 0
        pre_action = -1
        obs, _ = env.reset()
        hx = torch.zeros((net.lstm.num_layers, 1, net.lstm.hidden_size), dtype=torch.float32).to(device)
        cx = torch.zeros((net.lstm.num_layers, 1, net.lstm.hidden_size), dtype=torch.float32).to(device)
        while True:
            obs_v = ptan.agent.default_states_preprocessor([obs]).to(device)
            logits_v, _, (hx, cx) = net((obs_v, (hx, cx)))
            probs_v = F.softmax(logits_v, dim=1)
            probs = probs_v.data.cpu().numpy()
            action = np.argmax(probs)
            if action == 0 and pre_action == action:  # Noop
                noop_action_count += 1
                if noop_action_count > 30:
                    break
            else:
                noop_action_count = 0
            pre_action = action
            obs, reward, done, trunc, _ = env.step(action)
            total_reward += reward
            if done or trunc:
                break
    return total_reward / episodes


def optimized_states_preprocessor(states):
    """
    Convert list of states into the form suitable for model.
    :param states: list of numpy arrays with states
    :return: torch.Tensor
    """
    if len(states) == 1:
        np_states = np.expand_dims(states[0], 0)
    else:
        np_states = np.asarray([np.asarray(s) for s in states])
    return torch.from_numpy(np_states)


def select_device(args):
    if args.cuda and torch.cuda.is_available():
        return torch.device("cuda")
    elif torch.backends.mps.is_available() and args.cuda:
        return torch.device("mps")
    return torch.device("cpu")


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("--cuda", default=True, action="store_true", help="Enable cuda")
    parser.add_argument("-n", "--name", default="doubledunk", required=False, help="Name of the run")
    args = parser.parse_args()
    device = select_device(args)

    save_path = os.path.join("saves", "a2c-lstm-" + args.name)
    if not os.path.exists(save_path):
        os.makedirs(save_path)
    envs = [wrap_dqn(gym.make("ALE/Breakout-v5", obs_type='rgb', frameskip=4, repeat_action_probability=0.0), episodic_life=False) for _ in range(NUM_ENVS)]
    test_env = wrap_dqn(gym.make("ALE/Breakout-v5", obs_type='rgb', frameskip=4, repeat_action_probability=0.0), episodic_life=False)
    writer = SummaryWriter(comment="-a2c-lstm_" + args.name)

    net = AtariA2C(envs[0].observation_space.shape, envs[0].action_space.n).to(device)
    print(net)

    agent = DQNLstmAgent(net, NUM_ENVS, apply_softmax=True, device=device)
    exp_source = ptan.experience.A2CExperienceSourceDirect(envs, agent, gamma=GAMMA, steps_count=REWARD_STEPS)
    optimizer = optim.Adam(net.parameters(), lr=LEARNING_RATE, eps=1e-5, weight_decay=0.99)
    scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=50000, gamma=0.9)

    batch = []
    best_reward = 0
    frame_idx = 0
    train_count = 0

    if os.path.exists(save_path) and len(os.listdir(save_path)) > 0:
        # 增加加载模型的代码
        checkpoints = sorted(filter(lambda x: "epoch" in x, os.listdir(save_path)), key=lambda x: int(x.split('_')[2].split('.')[0]))
        checkpoint = torch.load(os.path.join(save_path, checkpoints[-1]), map_location=device, weights_only=False)
        frame_idx = checkpoint['frame_idx']
        train_count = checkpoint['train_count']
        net.load_state_dict(checkpoint['net'])
        optimizer.load_state_dict(checkpoint['optimizer'])
        scheduler.load_state_dict(checkpoint['scheduler'])
        print("加载模型成功")
        # 打印学习率
        for param_group in optimizer.param_groups:
            print("Learning rate:", param_group['lr'])

    with common.RewardTracker(writer, stop_reward=700) as tracker:
        with ptan.common.utils.TBMeanTracker(writer, batch_size=10) as tb_tracker:
            for step_idx, (mb_states, mb_rewards, mb_actions, mb_values, mb_hxs, mb_cxs) in enumerate(exp_source):

                new_rewards = exp_source.pop_total_rewards()
                if new_rewards:
                    if tracker.reward(np.mean(new_rewards), step_idx + frame_idx):
                        break


                states_v = torch.FloatTensor(mb_states).to(device)
                actions_v = torch.LongTensor(mb_actions).to(device)
                vals_ref_v = torch.FloatTensor(mb_rewards).to(device)
                hx_v = torch.FloatTensor(mb_hxs).to(device)
                cx_v = torch.FloatTensor(mb_cxs).to(device)
                pred_vals_v = torch.FloatTensor(mb_values).to(device)

                logits_v, value_v, _ = net((states_v, (hx_v, cx_v)))
                prob_v = F.softmax(logits_v, dim=1)
                log_prob_v = F.log_softmax(logits_v, dim=1)
                entropy_loss_v = -(prob_v * log_prob_v).sum(dim=1, keepdim=True)

                q_v = value_v.detach()
                
                policy_loss = 0
                value_loss = 0
                gae = torch.zeros(1, 1).to(device=device)

                for i in reversed(range(len(vals_ref_v) - 1)):
                    q_v = GAMMA * q_v + vals_ref_v[i]
                    adv_v = q_v - pred_vals_v[i]
                    value_loss = value_loss + VF_COEF * adv_v.pow(2)

                    delta_t = vals_ref_v[i] + GAMMA * pred_vals_v[i + 1] - pred_vals_v[i]
                    gae = gae * GAMMA * GAE_LAMBDA + delta_t

                    policy_loss = policy_loss - log_prob_v * gae.detach() - ENT_COEF * entropy_loss_v

                optimizer.zero_grad()
                (policy_loss.mean() + VF_COEF * value_loss.mean()).backward()
                nn_utils.clip_grad_norm_(net.parameters(), CLIP_GRAD)
                optimizer.step()
                train_count += 1
                # scheduler.step()

                if train_count % 200 == 0:
                    # Test the model
                    test_reward = test_model(test_env, net, device=device, episodes=2)
                    print(f"Test reward: {test_reward:.2f}")
                    common.save_best_model(test_reward, net.state_dict(), save_path, "a2c-lstm-best", keep_best=10)

                if train_count % SAVE_ITERS == 0:
                    checkpoint = {
                        "net": net.state_dict(),
                        "optimizer": optimizer.state_dict(),
                        "frame_idx": frame_idx + step_idx,
                        "scheduler": scheduler.state_dict(),
                        "train_count": train_count,
                    }
                    common.save_checkpoints(train_count, checkpoint, save_path, "a2c-lstm", keep_last=5)


                tb_tracker.track("advantage",       adv_v, step_idx + frame_idx)
                tb_tracker.track("values",          value_v, step_idx + frame_idx)
                tb_tracker.track("batch_rewards",   vals_ref_v, step_idx + frame_idx)
                tb_tracker.track("loss_entropy",    entropy_loss_v, step_idx + frame_idx)
                tb_tracker.track("loss_policy",     policy_loss.mean(), step_idx + frame_idx)
                tb_tracker.track("loss_value",      value_loss.mean(), step_idx + frame_idx)
                # tb_tracker.track("grad_l2",         np.sqrt(np.mean(np.square(grads))), step_idx + frame_idx)
                # tb_tracker.track("grad_max",        np.max(np.abs(grads)), step_idx + frame_idx)
                # tb_tracker.track("grad_var",        np.var(grads), step_idx + frame_idx)